package com.andnnl.list;

import java.util.*;
/**
 * Created by chenss on 2024/1/3.
 */
import java.util.HashSet;
import java.util.Set;

public class IPDomainWhitelistMatcher {

    private final TrieNode root;

    public IPDomainWhitelistMatcher(Set<String> ipDomainWhitelist) {
        root = new TrieNode();
        for (String ipOrDomain : ipDomainWhitelist) {
            insert(ipOrDomain);
        }
    }

    public boolean isWhitelisted(String ipOrDomain) {
        return search(ipOrDomain);
    }

    private void insert(String ipOrDomain) {
        TrieNode node = root;
        for (char c : ipOrDomain.toCharArray()) {
            if (c == '*') {
                node = insertWildcard(node);
            } else {
                if (!node.containsKey(c)) {
                    node.put(c, new TrieNode());
                }
                node = node.get(c);
            }
        }
        node.setEnd();
    }

    private TrieNode insertWildcard(TrieNode node) {
        TrieNode wildcardNode = new TrieNode();
        node.put('*', wildcardNode);
        return wildcardNode;
    }

    private boolean search(String ipOrDomain) {
        return searchHelper(ipOrDomain, 0, root);
    }

    private boolean searchHelper(String ipOrDomain, int index, TrieNode node) {
        if (index == ipOrDomain.length()) {
            return node.isEnd();
        }
        char c = ipOrDomain.charAt(index);
        if (c == '*') {
            Set<Character> keys = new HashSet<>();
            getAllKeysHelper(node, keys);
            for (char wildcardChar : keys) {
                if (searchHelper(ipOrDomain, index + 1, node.get(wildcardChar))) {
                    return true;
                }
            }
            return false;
        } else {
            if (!node.containsKey(c)) {
                return false;
            }
            return searchHelper(ipOrDomain, index + 1, node.get(c));
        }
    }

    private void getAllKeysHelper(TrieNode node, Set<Character> keys) {
        for (int i = 0; i< node.children.length; i++) {
            if (node.children[i] != null) {
                keys.add((char) i);
            }
        }
        if (node.containsKey('*')) {
            keys.add('*');
        }
    }

    public TrieNode getRoot() {
        return root;
    }

    @Override
    public String toString() {
        return "IPDomainWhitelistMatcher{" +
                "root=" + root +
                '}';
    }

    private static class TrieNode {
        private final TrieNode[] children;
        private boolean isEnd;

        public TrieNode() {
            children = new TrieNode[256];
            isEnd = false;
        }

        public void put(char c, TrieNode node) {
            children[c] = node;
        }

        public boolean containsKey(char c) {
            return children[c] != null;
        }

        public TrieNode get(char c) {
            return children[c];
        }

        public boolean isEnd() {
            return isEnd;
        }

        public void setEnd() {
            isEnd = true;
        }

        @Override
        public String toString() {
            return "TrieNode{" +
                    "children=" + Arrays.toString(children) +
                    ", isEnd=" + isEnd +
                    '}';
        }
    }

}